#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function

import numpy as np
import cv2
import copy

import almath
import motion
import vision_definitions

from numpy import deg2rad
from naoqi import ALProxy
from robot import tr2rpy,rpy2tr
from mdh_nao import Nao as mNao

# 不要用 from robot import * 因为该模块中大量使用 from xx import *
# 所以其内部命名空间很乱, 估计其中就有一个 copy 实际上是 numpy.copy
# 是一个函数,污染了我导入的 copy 模块

kin = mNao("nao")

class Nao(object):
    def __init__(self,ip=None,port=None):
        if (ip is not None) and (port is not None):
          self.ip = ip
          self.port = port
          self.motionProxy = ALProxy("ALMotion",ip,port)
          self.postureProxy = ALProxy("ALRobotPosture", ip, port)
          self.camProxy = ALProxy("ALVideoDevice",ip, port)
          self.frame = motion.FRAME_TORSO
          self.axisMask = almath.AXIS_MASK_ALL
          self.useSensorValues = True
          # resolution = vision_definitions.kVGA # 640x480
          resolution = vision_definitions.kQVGA # 320x240
          colorSpace = vision_definitions.kBGRColorSpace
          FPS = 25
          self.topClint = self.camProxy.subscribeCamera("top_cam_client",
                                                        vision_definitions.kTopCamera,
                                                        resolution,
                                                        colorSpace,
                                                        FPS)
          self.bottomClint = self.camProxy.subscribeCamera("bottom_cam_client",
                                                        vision_definitions.kBottomCamera,
                                                        resolution,
                                                        colorSpace,
                                                        FPS)
    def connect_red(self):
        self.__init__("192.168.1.106",9559) # 仿真
        #self.__init__("192.168.1.106",9559) # 真实

    def connect_blue(self):
        self.__init__("192.168.1.106",5960) # 仿真
        #self.__init__("192.168.1.106",9559) # 真实

    def __del__(self):
        if hasattr(self,"camProxy"):
            self.camProxy.unsubscribe(self.topClint)
            self.camProxy.unsubscribe(self.bottomClint)

    #### 视觉
    def get_top_image(self):
        return self._get_image(self.topClint)

    def get_bottom_image(self):
        return self._get_image(self.bottomClint)

    def _get_image(self,clint):
        naoImage = self.camProxy.getImageRemote(clint)
        if naoImage is None :
            if hasattr(self,"_image"):
                return self._image # 如果连接失败就返回缓存的图像
            else:
                return self._get_image() #
        else:
            imageWidth = naoImage[0]
            imageHeight = naoImage[1]
            array = np.frombuffer(naoImage[6],dtype=np.uint8)
            image = array.reshape((imageHeight, imageWidth,3))
            self._image = image # 图像缓存
            return image

    ##########   Motion 运动 ##########################
    # 笛卡尔空间
    def _goToPosture(self,postureName,v=1.0):
        self.postureProxy.goToPosture(postureName, v)

    def _getPos(self,name):
        Pos = self.motionProxy.getPosition(name,
                                         self.frame,
                                         self.useSensorValues)
        Pos = np.array(Pos)
        Pos[:3] *= 1000 # api长度单位为m,本系统为 mm
        return Pos

    def _setPos(self,name,Pos,v=1.0):
        assert isinstance(Pos,np.ndarray)
        Pos[:3] *= 0.001 # api长度单位为m,本系统为 mm
        Pos = Pos.tolist()
        self.motionProxy.setPositions(name,
                                      self.frame,
                                      Pos,
                                      v,
                                      self.axisMask)

    def _transformInterpolations(self,effectorList,pathList,axisMaskList,timeList):
        self.motionProxy.transformInterpolations(effectorList,
                                                self.frame,
                                                pathList,
                                                axisMaskList,
                                                timeList)
    def _positionInterpolations(self,effectorList,pathList,axisMaskList,timeList):
        self.motionProxy.positionInterpolations(effectorList,
                                                self.frame,
                                                pathList,
                                                axisMaskList,
                                                timeList)
    def _getTf(self,name):
        tf = self.motionProxy.getTransform(name,
                                          self.frame,
                                          self.useSensorValues)
        tf = np.array(tf).reshape((4,4))
        tf = np.matrix(tf)
        tf[0:3,3] *= 1000 # 默认长度单位为 mm
        return tf

    def _setTf(self,name,Tf,v):
        assert isinstance(Tf,np.matrix)
        Tf = np.array(Tf,dtype=np.float32)
        Tf[0:3,3] *= 0.001 # 默认长度单位为 mm, setTransforms 接受长度单位为 m
        Tf = Tf.flatten().tolist()
        self.motionProxy.setTransforms(name,
                                   self.frame,
                                   Tf,
                                   v,
                                   self.axisMask)

    # 关节空间
    def _setAngles(self,names,angles,fractionMaxSpeed=0.5):
        self.motionProxy.setAngles(names, angles, fractionMaxSpeed)

    def _getAngles(self,names):
        return self.motionProxy.getAngles(names, self.useSensorValues)

    # 电机电流(电机刚性)
    def _setStiffnesses(self,jointNames,stiffnesses):
        self.motionProxy.setStiffnesses(jointNames,stiffnesses)

    def _getStiffnesses(self,jointNames):
        return self.motionProxy.getStiffnesses(jointNames)

    # 手部控制
    def _openHand(self,handName):
        self.motionProxy.openHand(handName)

    def _closeHand(self,handName):
        self.motionProxy.closeHand(handName)

    # 定制操作
    def rest(self):
        self.motionProxy.rest()

    def wakeUp(self):
        self.motionProxy.wakeUp()

    def isWakeUp(self):
        return self.motionProxy.robotIsWakeUp()

    def Stand(self):
        self._goToPosture("Stand") # 完全站立起来

    def StandInit(self):
        self._goToPosture("StandInit") # wakeUp 后的姿态

    def StandZero(self):
        self._goToPosture("StandZero") # 仿真开始时的姿态

    def Crouch(self):
        self._goToPosture("Crouch") # rest 后的姿态

    def getLArmPos(self):
        return self._getPos("LArm")

    def setLArmPos(self,Pos,v=0.5):
        self._setPos("LArm",Pos,v)

    def getRArmPos(self):
        return self._getPos("RArm")

    def setRArmPos(self,Pos,v=0.5):
        self._setPos("RArm",Pos,v)

    def getLArmTf(self):
        return self._getTf("LArm")

    def setLArmTf(self,Tf,v=0.5):
        self._setTf("LArm",Tf,v)

    def getRArmTf(self):
        return self._getTf("RArm")

    def setRArmTf(self,Tf,v=0.5):
        self._setTf("RArm",Tf,v)

    def getHeadAngles(self):
        return self._getAngles("Head")

    def setHeadAngles(self,angels):
        self._setAngles("Head",angels)

    def getLArmAngles(self):
        return self._getAngles("LArm")

    def setLArmAngles(self,angels):
        self._setAngles("LArm",angels)

    def getRArmAngles(self):
        return self._getAngles("RArm")

    def setRArmAngles(self,angels):
        self._setAngles("RArm",angels)

    def mirrorRArm(self,Langles=None): # 将RArm放置为LArm的镜像位姿
        angles = self.getLArmAngles() if Langles is None else Langles
        self.setRArmAngles(self.mirrorArmAngles(angles))

    def mirrorArmAngles(self,angles):
        mangles = copy.copy(angles)
        mangles[1]*=-1
        mangles[2]*=-1
        mangles[3]*=-1
        mangles[4]*=-1
        return mangles

    def mirrorArmPose(self,Pos):
        mPos = Pos.copy()
        mPos.y *= -1
        mPos.wx *= -1
        mPos.wz *= -1
        return mPos

    def closeLHand(self):
        self._closeHand("LHand")

    def openLHand(self):
        self._openHand("LHand")

    def closeRHand(self):
        self._closeHand("RHand")

    def openRHand(self):
        self._openHand("RHand")

    def tf2Pos(self,Tf):
        x,y,z = Tf[:3,3]
        wz,wy,wx = np.array(tr2rpy(Tf)).flatten()
        return np.array([x,y,z,wx,wy,wz])

    def pos2Tf(self,Pos):
        x,y,z,wx,wy,wz = Pos
        Tf = np.matrix(np.eye(4))
        Tf[:3,3] = np.matrix([x,y,z]).transpose()
        Tf[:3,:3] = rpy2tr(wz,wy,wx)
        return Tf

    def tf2Rodrigues(self,Tf):
        rodrigues, jaco = cv2.Rodrigues(Tf[:3,:3])
        return rodrigues

    def rodrigues2Tf(self,rodrigues):
        R,jaco = cv2.Rodrigues(rodrigues)
        Tf = np.eye(4,dtype=R.dtype)
        Tf[:3,:3] = R
        return Tf

    def rodrigues2Pos(self,rodrigues):
        Tf = self.rodrigues2Tf(rodrigues)
        return self.tf2Pos(Tf)

    def inv(self,tp):
        if isinstance(tp,almath.Transform):
            return tp.inverse()
        elif isinstance(tp,(np.mat,np.ndarray,np.matrix)):
            return np.linalg.inv(tp)

    def fLeftHand(self,th1,th2,th3,th4,th5,th6):
        q = [th1,th2,th3,th4,th5]
        return kin.LeftHand.fkine(q)

    def iLeftHand(self,Tf):
        assert isinstance(Tf,np.matrix)
        q,err = kin.LeftHand.ikine(Tf) # 解析解
        if err > 1.0:
            q0 = self.getLArmAngles()
            q0 = q0[0:5]
            q = kin.LeftHand.ikine2(Tf,q0,(0,1,2,3,5)) # 数值解, 保持位置和末端y轴不动
        th1,th2,th3,th4,th5 = np.array(q).flatten()
        return th1,th2,th3,th4,th5, 0

    def setLArmTf2(self,Tf):
        angles = self.iLeftHand(Tf)
        self.setLArmAngles(angles)

    def fRightHand(self,th1,th2,th3,th4,th5,th6):
        q = [th1,th2,th3,th4,th5]
        return kin.RightHand.fkine(q)

    def iRightHand(self,Tf):
        assert isinstance(Tf,np.matrix)
        q,err = kin.RightHand.ikine(Tf) # 解析解
        if err > 1.0:
            q0 = self.getRArmAngles()
            q0 = q0[0:5]
            q = kin.RightHand.ikine2(Tf,q0,(0,1,2,3,5)) # 数值解, 保持位置和末端z轴不动
        th1,th2,th3,th4,th5 = np.array(q).flatten()
        return th1,th2,th3,th4,th5, 0

    def setRArmTf2(self,Tf):
        angles = self.iRightHand(Tf)
        self.setRArmAngles(angles)

    def simulationInitPose(self):
        if self.isWakeUp():
            self.StandInit()
        else:
            self.wakeUp()
        self.setHeadAngles([0,deg2rad(-13.3)]) # 头部位姿
        self.closeLHand() # 握紧手指
        self.closeRHand() # 握紧手指
        # 指定姿态
        self.frame=motion.FRAME_TORSO
        self.axisMask=motion.AXIS_MASK_ALL
        self.useSensorValues=True
        #Pos = almath.Position6D(0.180455, 0.076655, 0.0261944,
        #                        -1.56787, -0.000376985, -0.257331)
        angles = [0.786008536816,0.046681009233,-1.17838895321,-0.83504575491,-0.29978209734,0.0]
        self.setLArmAngles(angles)
        self.setRArmAngles(self.mirrorArmAngles(angles))


